import warnings
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from statsmodels.tsa.stattools import adfuller
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf
from statsmodels.tsa.arima.model import ARIMA
from statsmodels.tsa.statespace.sarimax import SARIMAX
import pmdarima as pm
from pmdarima.model_selection import train_test_split
from sklearn.metrics import mean_squared_error
import pdb

warnings.filterwarnings('ignore')


def read_data():
    file = 'PM2.5.xlsx'
    df = pd.read_excel(file)# index_col=['Time','PM2.5'], parse_dates=['Time']
    # print("done")
    # print(df)

    ts = pd.to_datetime(df['Time'], format='%Y-%m-%d')

    # ts = df['2015PM']

    return ts


def plot(ts):
    results = adfuller(ts)
    results_str = 'ADF test, p-value is: {}'.format(results[1])

    grid = plt.GridSpec(2, 2)
    ax1 = plt.subplot(grid[0, :])
    ax2 = plt.subplot(grid[1, 0])
    ax3 = plt.subplot(grid[1, 1])

    ax1.plot(ts)
    ax1.set_title(results_str)
    plot_acf(ts, lags=int(len(ts) / 2 - 1), ax=ax2)
    plot_pacf(ts, lags=int(len(ts) / 2 - 1), ax=ax3)
    plt.show()


ts = read_data()
print(ts)
plot(ts)


def find_pq(ts, d=0, max_p=5, max_q=5):
    best_p, best_q = 0, 0
    best_aic = np.inf

    for p in range(max_p):
        for q in range(max_q):
            model = ARIMA(ts, order=(p, d, q)).fit()
            aic = model.aic

            if aic < best_aic:
                best_aic = aic
                best_p = p
                best_q = q

    return best_p, best_q, best_aic


def version_arima_with_manual(ts):
    """
    ARIMA（手动季节差分）
    """
    # 周期大小
    periods = 365

    # 季节差分
    ts_diff = ts - ts.shift(periods)
    # 再次差分（季节差分后p值小于0.05-接近，可认为平稳，若要严格一点也可再做一次差分）
    # ts_diff = ts_diff - ts_diff.shift(1)

    # （训练数据中不能有缺失值，这里差分后前几个值为nan，故去除）
    ts_diff = ts_diff[~pd.isnull(ts_diff)]

    # 数据拆分
    train, test = train_test_split(ts_diff, train_size=0.8)

    # 模型训练（训练数据为差分后的数据-已平稳，所以d=0）
    p, q, _ = find_pq(train)
    model = ARIMA(train, order=(p, 0, q)).fit()
    print(model.summary())

    # 拟合结果
    fitted = model.fittedvalues

    # 模型预测
    fcst = model.forecast(test.shape[0])

    # 差分还原（拟合结果）
    fitted += ts.shift(periods)

    # 差分还原（预测结果）
    tmp = ts.loc[train.index].values.tolist() + fcst.values.tolist()
    for i in range(len(tmp) - fcst.shape[0], len(tmp)):
        tmp[i] += tmp[i - periods]
    fcst.loc[:] = tmp[-fcst.shape[0]:]

    # 模型评估
    rmse = np.sqrt(mean_squared_error(test, fcst))
    print('RMSE: %.4f' % rmse)

    # 可视化
    plt.figure(figsize=(12, 4))
    ts.plot(label='Ads')
    fitted.plot(label='fitted')
    fcst.plot(label='forecast')
    plt.legend()
    plt.grid(True)
    plt.show()

print('done')
version_arima_with_manual(ts)